import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
from sklearn import datasets
# from pandas import DataFrame
# import pandas as pd

# 1.准备数据
## 1.1 数据集读入
x_data = datasets.load_iris().data      #数据
y_data = datasets.load_iris().target    #标签

## 1.2 数据集乱序
### 通过相同seed保证标签和特征打乱的顺序一致
np.random.seed(116)
np.random.shuffle(x_data)
np.random.seed(116)
np.random.shuffle(y_data)
tf.random.set_seed(116)

## 1.3 数据集分出（永不相见的）训练集和测试集
### 数据切片
x_train = x_data[:-30]  #从第一个元素到倒数第30个元素之前的所有元素
y_train = y_data[:-30]
x_test = x_train[-30:]  #从序列的倒数第30个元素开始，一直到序列的末尾
y_test = y_train[-30:]

### 将数据类型更新为float32
x_train = tf.cast(x_train, tf.float32)
x_test = tf.cast(x_test, tf.float32)

## 1.4 配对（输入特征、标签），每次喂一小撮（batch）
### 数据集（训练集、测试集）切片（每32个输入为一个批量大小）
### 注意批量大小（batch size）对训练的影响！！！
train_db = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(32)
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)


# 2.搭建网络
w1 = tf.Variable(tf.random.truncated_normal([4,3], stddev=0.1, seed=1))
b1 = tf.Variable(tf.random.truncated_normal([3], stddev=0.1, seed=1))

lr = 0.1                #学习率:过大的学习率可能导致优化过程不稳定甚至发散，而过小的学习率则会导致收敛速度缓慢
train_loss_results = []
test_acc = []
epoch = 500             #迭代次数, 建议引入EarlyStopping，当loss不再下降时可以减少消耗
loss_all = 0
trained_cnt = 0

## 训练部分
for epoch in range(epoch): #迭代500次
    for step, (x_train, y_train) in enumerate(train_db):    # 使用 enumerate 获取批次索引（step)和批次数据(x_train, y_train)
        with tf.GradientTape() as tape:
            y = tf.matmul(x_train, w1) + b1     #计算得分
            y = tf.nn.softmax(y)                #将得分转化为概率（激活函数）
            y_ = tf.one_hot(y_train, depth=3)   #将类别标签转换为向量形式（独热编码）
            # 计算均方误差
            loss = tf.reduce_mean(tf.square(y_ - y))
            loss_all += loss.numpy()
        grads = tape.gradient(loss, [w1, b1]) #计算梯度
        w1.assign_sub(lr * grads[0])            #使用梯度下降优化(下降值为lr * grads[0])，更新权重w1
        b1.assign_sub(lr * grads[1])
        # print("w1: ", w1)
        # print("b1: ", b1)
    print("Epoch {}, loss: {}".format(epoch, loss_all/4))
    train_loss_results.append(loss_all / 4)
    loss_all = 0

## 测试部分
total_correct, total_number = 0, 0
for x_test, y_test in test_db:
    y = tf.matmul(x_test, w1) + b1
    y = tf.nn.softmax(y)
    pred = tf.argmax(y, axis=1)
    pred = tf.cast(pred, dtype=y_test.dtype)
    correct = tf.cast(tf.equal(pred, y_test), dtype=tf.int32)
    correct = tf.reduce_sum(correct)
    total_correct += int(correct)
    total_number += x_test.shape[0]
acc = total_correct / total_number
test_acc.append(acc)
print("Test acc:", acc)
print("-------------------------")

# 3.acc/loss 可视化
## 绘制loss曲线
plt.title('Loss Function Curve')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.plot(train_loss_results, label="$Loss$")
plt.legend()
plt.show()

## 绘制Accuracy曲线
plt.title('Acc Curve')
plt.xlabel('Epoch')
plt.ylabel('Acc')
plt.plot(test_acc, label="$Accuracy$")
plt.legend()
plt.show()

